'''压缩CG包使用的脚本.压缩相似的cg.
  指定一个CG包目录,相似的cg文件的文件名须有连续性.
  主类为CGCompresser,其中:
  compareCG方法用于比较CG包的图片,制作压缩任务;
  compress方法用于执行压缩;
  main方法一键执行.
  第三方压缩命令(这里用7z),以及图像的相似度定义在主类的构造方法中传入.
  
  使用者在开头的CG_PATH,OUTPUT_PATH指定CG目录与输出目录,
  _7Z_EXE指定7z压缩程序的位置.
  
  经过实验,最好的是用7z的lzma压缩算法,压缩率高,支持的看图软件也多一点,
  Hamana就支持.LZMA2压缩算法支持的看图软件少些,我用的是Hamana,不支持.
  
  ---
  Last update @2019/8/15
  By setycyas
'''

## 在这里指定CG包路径与压缩后的路径
CG_PATH = r'F:\CG\PC官方\冥色の隷姫～緩やかに廃滅する青珊瑚の森～'
OUTPUT_PATH = r'F:\CG\PC官方\冥色の隷姫_lzma'
_7Z_EXE = r'D:\software\7-Zip\7z.exe'

import os
import sys
import shutil
import subprocess

from cv2 import cv2
import numpy as np

class CGCompresser:

  ## 初始化参数:CG包目录,输出目录,定义最低相似度(相似度高于这个数认为是相似图),
  ## 压缩命令表orderList, 是压缩的前面参数,后面的参数是输出文件与压缩文件,
  ## 使用时再把参数补充完整.
  def __init__(self, cgPath, oPath, similar = 0.3,
    orderList = [_7Z_EXE, '-t7z', '-m0=LZMA', 'a']):
    ## 参数复制
    self.cgPath = cgPath
    assert(os.path.isdir(cgPath))
    print(r'源CG目录: {}'.format(self.cgPath))
    self.oPath = oPath
    if os.path.isdir(self.oPath):
      print(r'目标目录: {}'.format(self.oPath))
    else:
      os.mkdir(self.oPath)
      print(r'目标目录: {} 新建成功!'.format(self.oPath))
    self.similar = similar
    self.orderList = orderList
    ## 压缩任务表,每一项为需压缩的文件列表
    self.missionList = []
    return
    
  ## 判断一个文件对象是否图片文件,是则返回True,否则返回False
  def _isPicture(self, filepath):
    ## 图片扩展名表,小写
    picExtNames = ['.bmp', '.jpg', '.jpeg', '.gif', '.png']
    extName = os.path.splitext(filepath)[1].lower()
    if extName in picExtNames:
      return True
    else:
      return False
    
  ## 对cg包的文件进行分析,制作压缩任务表.
  def compareCG(self):
    ## 制作有效的图片文件列表fileList,含目录
    fileList = os.listdir(self.cgPath)
    fileList.sort()
    fileList = list(filter(self._isPicture, fileList))
    fileList = [os.path.join(self.cgPath, filename) for filename in fileList]
    ## 按顺序比较图片文件,制作任务表
    curFile = fileList[0] # 当前代表图形文件
    curImg = cv2.imdecode(np.fromfile(curFile, dtype = np.uint8), -1) 
    curMission = [curFile] # 当前正在生成的压缩任务表
    fileNumber = len(fileList)
    print('\n图片分析进度:')
    for i in range(1, fileNumber):
      sys.stdout.write('{0} / {1} \r'.format(i, fileNumber))
      sys.stdout.flush()
      filename = fileList[i]
      img = cv2.imdecode(np.fromfile(filename, dtype = np.uint8), -1)
      imgCmp = (curImg.shape == img.shape)
      imgCmp = imgCmp and ((curImg == img).mean() > self.similar)
      ## 比较当前代表图与后续图片,相似的加入同类任务表,不相似则保存当前任务,
      ## 并建立新任务
      if imgCmp:
        curMission.append(filename)
      else:
        self.missionList.append(curMission)
        curFile = filename
        curImg = img
        curMission = [filename]
    self.missionList.append(curMission)
    print('图片分析完成!\n')
    return
      
  ## 压缩任务表变成一个字符串,以便用于更多记录显示.测试分析用
  def _missions2str(self):
    result = []
    for mission in self.missionList:
      result.append('='*20)
      for filename in mission:
        result.append(os.path.basename(filename))
    result = '\n'.join(result)
    return result
    
  ## 压缩所有任务
  def compress(self):
    print('开始调用7z压缩图片:')
    for i in range(len(self.missionList)):
      mission = self.missionList[i]
      ## 找到首图片,以首图片命名压缩包.制作orderList以便调用第三方软件压缩.
      topFile = mission[0]
      ## 如果任务表只有一个图片,则不进行压缩与缩略图制作,直接把唯一图片复制过去
      if (len(mission) == 1):
        copyFile = os.path.join(self.oPath, os.path.basename(topFile))
        shutil.copy(topFile, copyFile)
        continue
      ## 任务表超过一个图片,则生成缩略图,然后调用7z进行压缩
      ## 缩略图生成
      thumbnailFile = os.path.join(self.oPath, os.path.basename(topFile))
      img = cv2.imdecode(np.fromfile(topFile, dtype = np.uint8), -1)
      img = cv2.resize(img, (img.shape[1] // 4, img.shape[0] // 4))
      ## 以下这一段是保存修改后的img到缩略图,抄过来的代码.
      with open(thumbnailFile, 'wb') as f:
        if img.dtype in (np.float64, np.float32, np.float):
            img = (img * 256).astype(np.uint8)       
        _, data = cv2.imencode(thumbnailFile, img)
        f.write(data)  
      ## 调用7z进行压缩
      orderList = self.orderList.copy()  
      orderList.append(
        (os.path.join(self.oPath, os.path.basename(topFile))+'.7z')
      )
      for j in range(0, len(mission)):
        filename = mission[j]
        orderList.append(filename)
      subprocess.call(orderList)
    ## 压缩任务全部完成      
    print('压缩任务全部完成!')
    return
    
  ## 主函数
  def main(self):
    self.compareCG()
    self.compress()
    return
  
  
if __name__ == '__main__':
  ## 输入输出目录
  cgPath =  CG_PATH
  oPath = OUTPUT_PATH
  ## 确定运行
  x = input('CG压缩任务即将开始,可能耗时很多,输入1回车开始,其他结束')
  if (x == '1'):
    cgc = CGCompresser(cgPath, oPath)
    cgc.main()
  else:
    print('退出任务')
  